import os
import jwt
import datetime
from functools import wraps
from werkzeug.security import generate_password_hash, check_password_hash
import pandas as pd
import numpy as np
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from io import BytesIO
import base64
from flask import Flask, render_template, request, redirect, url_for, flash, send_file, jsonify, session
from flask_sqlalchemy import SQLAlchemy
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-super-secret-key-keep-it-safe'  # 生产环境请更换
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MODEL_FOLDER'] = 'saved_models'
app.config['PLOTS_FOLDER'] = 'static/plots'
app.config['RESULT_FOLDER'] = 'results'
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///users.db'  # SQLite数据库
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

# 初始化数据库
db = SQLAlchemy(app)

# 用户模型
class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(50), unique=True, nullable=False)
    email = db.Column(db.String(100), unique=True, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.datetime.utcnow)

    def set_password(self, password):
        self.password_hash = generate_password_hash(password)

    def check_password(self, password):
        return check_password_hash(self.password_hash, password)

# 创建数据库表（首次运行时需要）
with app.app_context():
    db.create_all()

# 登录装饰器 - 保护需要登录的路由
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user_id' not in session:
            flash('请先登录', 'warning')
            return redirect(url_for('login', next=request.url))
        return f(*args, **kwargs)
    return decorated_function

# 确保必要的文件夹存在
for folder in [app.config['UPLOAD_FOLDER'], app.config['MODEL_FOLDER'], 
               app.config['PLOTS_FOLDER'], app.config['RESULT_FOLDER']]:
    if not os.path.exists(folder):
        os.makedirs(folder)

# 模型文件路径
MODEL_PATH = os.path.join(app.config['MODEL_FOLDER'], 'trained_model.pkl')
FEATURES_PATH = os.path.join(app.config['MODEL_FOLDER'], 'features.pkl')
RESULT_PATH = os.path.join(app.config['RESULT_FOLDER'], 'prediction_results.csv')

# 配置matplotlib支持中文
plt.rcParams["font.family"] = ["Noto Sans SC", "SimHei", "WenQuanYi Micro Hei"]
plt.rcParams['axes.unicode_minus'] = False

# 注册路由
@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form['username']
        email = request.form['email']
        password = request.form['password']
        confirm_password = request.form['confirm_password']

        # 验证表单
        if not username or not email or not password:
            flash('所有字段都是必填的', 'danger')
            return redirect(url_for('register'))
            
        if password != confirm_password:
            flash('两次密码输入不一致', 'danger')
            return redirect(url_for('register'))
            
        # 检查用户名和邮箱是否已存在
        if User.query.filter_by(username=username).first():
            flash('用户名已被注册', 'danger')
            return redirect(url_for('register'))
            
        if User.query.filter_by(email=email).first():
            flash('邮箱已被注册', 'danger')
            return redirect(url_for('register'))
            
        # 创建新用户
        new_user = User(username=username, email=email)
        new_user.set_password(password)
        
        try:
            db.session.add(new_user)
            db.session.commit()
            flash('注册成功，请登录', 'success')
            return redirect(url_for('login'))
        except Exception as e:
            db.session.rollback()
            flash(f'注册失败: {str(e)}', 'danger')
            return redirect(url_for('register'))
            
    return render_template('register.html')

# 登录路由
@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        
        user = User.query.filter_by(username=username).first()
        
        if not user or not user.check_password(password):
            flash('用户名或密码错误', 'danger')
            return redirect(url_for('login'))
            
        # 登录成功，设置会话
        session['user_id'] = user.id
        session['username'] = user.username
        
        # 重定向到之前请求的页面或首页
        next_page = request.args.get('next')
        return redirect(next_page or url_for('index'))
        
    return render_template('login.html')

# 登出路由
@app.route('/logout')
def logout():
    session.pop('user_id', None)
    session.pop('username', None)
    flash('已成功登出', 'success')
    return redirect(url_for('index'))

# 首页路由
@app.route('/')
def index():
    model_exists = os.path.exists(MODEL_PATH)
    return render_template('index.html', 
                          model_exists=model_exists,
                          current_user=session.get('username'))

# 训练路由
@app.route('/train', methods=['GET', 'POST'])
@login_required
def train():
    model_exists = os.path.exists(MODEL_PATH)
    
    if request.method == 'POST':
        if 'file' not in request.files:
            flash('请上传训练数据文件', 'danger')
            return redirect(request.url)
        
        file = request.files['file']
        if file.filename == '':
            flash('未选择文件', 'danger')
            return redirect(request.url)
        
        if file and file.filename.endswith('.csv'):
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], 'training_data.csv')
            file.save(file_path)
            
            try:
                df = pd.read_csv(file_path)
                
                if 'paid_seat' not in df.columns:
                    flash('训练数据必须包含"paid_seat"标签列', 'danger')
                    return redirect(url_for('train'))
                
                X = df.drop('paid_seat', axis=1)
                y = df['paid_seat']
                
                # 保存特征列表
                features = X.columns.tolist()
                joblib.dump(features, FEATURES_PATH)
                
                # 划分训练集和测试集
                X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
                
                # 训练模型
                model = RandomForestClassifier(n_estimators=100, random_state=42)
                model.fit(X_train, y_train)
                
                # 评估模型
                y_pred = model.predict(X_test)
                accuracy = accuracy_score(y_test, y_pred)
                
                # 保存模型
                joblib.dump(model, MODEL_PATH)
                
                # 生成评估图表
                generate_evaluation_plots(y_test, y_pred)
                
                flash(f'模型训练完成！准确率: {accuracy:.2f}', 'success')
                return redirect(url_for('index'))
                
            except Exception as e:
                flash(f'训练过程出错: {str(e)}', 'danger')
                return redirect(url_for('train'))
    
    return render_template('train.html', model_exists=model_exists, current_user=session.get('username'))

# 删除模型路由
@app.route('/delete_model', methods=['POST'])
@login_required
def delete_model():
    try:
        if os.path.exists(MODEL_PATH):
            os.remove(MODEL_PATH)
        if os.path.exists(FEATURES_PATH):
            os.remove(FEATURES_PATH)
        flash('模型已成功删除', 'success')
    except Exception as e:
        flash(f'删除模型失败: {str(e)}', 'danger')
    return redirect(url_for('index'))

# 预测路由
@app.route('/predict', methods=['GET', 'POST'])
@login_required
def predict():
    if not os.path.exists(MODEL_PATH) or not os.path.exists(FEATURES_PATH):
        flash('请先训练模型', 'warning')
        return redirect(url_for('train'))
    
    prediction_done = False
    result_df = None
    plots = {}
    
    if request.method == 'POST':
        if 'file' not in request.files:
            flash('请上传预测数据文件', 'danger')
            return redirect(request.url)
        
        file = request.files['file']
        if file.filename == '':
            flash('未选择文件', 'danger')
            return redirect(request.url)
        
        if file and file.filename.endswith('.csv'):
            try:
                model = joblib.load(MODEL_PATH)
                train_features = joblib.load(FEATURES_PATH)
                
                df = pd.read_csv(file.stream)
                predict_df = align_features(df, train_features)
                
                predictions = model.predict(predict_df)
                predict_proba = model.predict_proba(predict_df)[:, 1]
                
                result_df = df.copy()
                result_df['paid_seat_prediction'] = predictions
                result_df['paid_seat_probability'] = predict_proba
                
                result_df.to_csv(RESULT_PATH, index=False)
                plots = generate_prediction_plots(result_df)
                prediction_done = True
                
            except Exception as e:
                flash(f'预测过程出错: {str(e)}', 'danger')
                return redirect(url_for('predict'))
    
    return render_template('predict.html', 
                          prediction_done=prediction_done, 
                          result_df=result_df,
                          plots=plots,
                          current_user=session.get('username'))

# 下载结果路由
@app.route('/download_result')
@login_required
def download_result():
    if os.path.exists(RESULT_PATH):
        return send_file(RESULT_PATH, as_attachment=True)
    else:
        flash('没有可下载的预测结果', 'danger')
        return redirect(url_for('predict'))

# 分析路由
@app.route('/analysis')
@login_required
def analysis():
    train_data_path = os.path.join(app.config['UPLOAD_FOLDER'], 'training_data.csv')
    has_data = os.path.exists(train_data_path)
    plots = {}
    
    if has_data:
        try:
            df = pd.read_csv(train_data_path)
            plots = generate_analysis_plots(df)
        except Exception as e:
            flash(f'数据分析出错: {str(e)}', 'danger')
    
    return render_template('analysis.html', 
                          has_data=has_data, 
                          plots=plots,
                          current_user=session.get('username'))

# 辅助函数：特征对齐
def align_features(predict_df, train_features):
    missing_features = set(train_features) - set(predict_df.columns)
    for feature in missing_features:
        predict_df[feature] = 0
    
    extra_features = set(predict_df.columns) - set(train_features)
    predict_df = predict_df.drop(columns=extra_features)
    
    return predict_df[train_features]

# 辅助函数：生成评估图表
def generate_evaluation_plots(y_test, y_pred):
    cm = confusion_matrix(y_test, y_pred)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=['不付费', '付费'], 
                yticklabels=['不付费', '付费'])
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    plt.title('混淆矩阵')
    cm_path = os.path.join(app.config['PLOTS_FOLDER'], 'confusion_matrix.png')
    plt.savefig(cm_path)
    plt.close()

# 辅助函数：生成预测结果图表
def generate_prediction_plots(result_df):
    plots = {}
    
    plt.figure(figsize=(8, 6))
    sns.countplot(x='paid_seat_prediction', data=result_df)
    plt.xlabel('预测结果')
    plt.ylabel('数量')
    plt.title('付费选座预测分布')
    plt.xticks([0, 1], ['不付费', '付费'])
    
    buffer = BytesIO()
    plt.savefig(buffer, format='png')
    plots['prediction_distribution'] = base64.b64encode(buffer.getvalue()).decode('utf-8')
    plt.close()
    
    plt.figure(figsize=(8, 6))
    sns.histplot(result_df['paid_seat_probability'], bins=20, kde=True)
    plt.xlabel('付费概率')
    plt.ylabel('数量')
    plt.title('付费选座概率分布')
    
    buffer = BytesIO()
    plt.savefig(buffer, format='png')
    plots['probability_distribution'] = base64.b64encode(buffer.getvalue()).decode('utf-8')
    plt.close()
    
    return plots

# 辅助函数：生成分析图表
def generate_analysis_plots(df):
    plots = {}
    
    plt.figure(figsize=(8, 6))
    sns.barplot(x='member_level', y='paid_seat', data=df)
    plt.xlabel('会员等级 (0=普通, 3=白金卡)')
    plt.ylabel('付费选座比例')
    plt.title('会员等级与付费选座意愿')
    
    buffer = BytesIO()
    plt.savefig(buffer, format='png')
    plots['member_level'] = base64.b64encode(buffer.getvalue()).decode('utf-8')
    plt.close()
    
    plt.figure(figsize=(8, 6))
    sns.barplot(x='trip_purpose', y='paid_seat', data=df)
    plt.xlabel('出行目的 (0=商务, 1=旅游, 2=探亲, 3=其他)')
    plt.ylabel('付费选座比例')
    plt.title('出行目的与付费选座意愿')
    
    buffer = BytesIO()
    plt.savefig(buffer, format='png')
    plots['trip_purpose'] = base64.b64encode(buffer.getvalue()).decode('utf-8')
    plt.close()
    
    plt.figure(figsize=(8, 6))
    sns.barplot(x='route_type', y='paid_seat', data=df)
    plt.xlabel('航线类型 (0=短途, 1=中途, 2=长途)')
    plt.ylabel('付费选座比例')
    plt.title('航线类型与付费选座意愿')
    
    buffer = BytesIO()
    plt.savefig(buffer, format='png')
    plots['route_type'] = base64.b64encode(buffer.getvalue()).decode('utf-8')
    plt.close()
    
    return plots

if __name__ == '__main__':
    app.run(debug=True)